﻿#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import time
import torch
import torchaudio
import numpy as np
from flask import Flask, request, send_file, jsonify
from scipy.io import wavfile
import tempfile
import logging
from cosyvoice.cli.cosyvoice import CosyVoice2
from cosyvoice.utils.file_utils import load_wav

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# 配置参数
REFERENCE_AUDIO_PATH = "./asset/zbs.WAV"
REFERENCE_TEXT = "以后啊咱们爷们都别喝了对不对，警察呀多不容易啊，这一年四季，春夏秋冬，站在油漆路上，就搁那指挥交通。"

# 全局模型和参考音频
cosyvoice = None
reference_audio = None

def init_tts_system():
    """初始化文本转语音系统"""
    global cosyvoice, reference_audio
    
    logger.info("正在加载文本转语音模型...")
    try:
        # 加载CosyVoice2模型
        cosyvoice = CosyVoice2(
            'pretrained_models/CosyVoice2-0.5B',
            load_jit=False,
            load_trt=False,
            load_vllm=False,
            fp16=False
        )
        
        # 加载参考音频
        reference_audio = load_wav(REFERENCE_AUDIO_PATH, 16000)
        logger.info(f"文本转语音模型加载完成，参考音频: {REFERENCE_AUDIO_PATH}")
        logger.info(f"参考音频长度: {len(reference_audio)/16000:.2f}秒")
        return True
    except Exception as e:
        logger.error(f"加载文本转语音模型失败: {str(e)}")
        return False

def generate_speech(text):
    """生成语音的核心函数"""
    try:
        logger.info(f"开始合成语音: {text[:50]}...")
        
        # 生成语音
        results = list(cosyvoice.inference_zero_shot(
            text,  # 要合成的文本
            REFERENCE_TEXT,  # 参考音频对应的文本
            reference_audio,  # 参考音频
            stream=False,
            text_frontend=False
        ))
        
        if not results:
            logger.error("未生成语音结果")
            return None, None
            
        # 处理结果
        speech_data = results[0]['tts_speech']
        sample_rate = cosyvoice.sample_rate
        
        # 确保数据类型正确
        if speech_data.dtype != torch.float32:
            speech_data = speech_data.to(torch.float32)
        
        # 转换为numpy数组
        speech_np = speech_data.cpu().numpy()
        
        # 归一化音频
        max_val = np.max(np.abs(speech_np))
        if max_val > 0:
            speech_np = speech_np / max_val
        
        return speech_np, sample_rate
        
    except Exception as e:
        logger.error(f"语音生成失败: {str(e)}")
        return None, None

@app.route('/tts', methods=['POST'])
def text_to_speech():
    """文本转语音API端点"""
    start_time = time.time()
    
    # 获取请求数据
    data = request.json
    text = data.get('text')
    
    if not text:
        return jsonify({"error": "缺少'text'参数"}), 400
    
    # 生成语音
    speech_np, sample_rate = generate_speech(text)
    
    if speech_np is None or sample_rate is None:
        return jsonify({"error": "语音生成失败"}), 500
    
    try:
        # 创建临时文件
        with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file:
            temp_filename = temp_file.name
        
        # 保存为标准的16位PCM WAV文件
        wavfile.write(temp_filename, sample_rate, (speech_np * 32767).astype(np.int16))
        
        # 记录音频信息
        duration = len(speech_np) / sample_rate
        logger.info(f"语音生成成功 - 时长: {duration:.2f}秒 - 文本长度: {len(text)}")
        
        return send_file(
            temp_filename, 
            mimetype='audio/wav',
            as_attachment=False,
            download_name='speech.wav'
        )
        
    except Exception as e:
        logger.error(f"文件保存失败: {str(e)}")
        return jsonify({"error": str(e)}), 500
    finally:
        # 稍后删除临时文件
        try:
            if temp_filename and os.path.exists(temp_filename):
                os.unlink(temp_filename)
        except:
            pass

@app.route('/health', methods=['GET'])
def health_check():
    """健康检查端点"""
    return jsonify({
        "status": "运行正常", 
        "model_loaded": cosyvoice is not None,
        "reference_audio": REFERENCE_AUDIO_PATH,
        "reference_text": REFERENCE_TEXT
    })

@app.route('/test', methods=['GET'])
def test_tts():
    """测试端点"""
    test_text = "这是一个测试文本，用于验证语音合成功能是否正常工作。"
    speech_np, sample_rate = generate_speech(test_text)
    
    if speech_np is None or sample_rate is None:
        return jsonify({"error": "测试语音生成失败"}), 500
    
    try:
        # 创建临时文件
        with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file:
            temp_filename = temp_file.name
        
        # 保存为WAV文件
        wavfile.write(temp_filename, sample_rate, (speech_np * 32767).astype(np.int16))
        
        return send_file(
            temp_filename, 
            mimetype='audio/wav',
            as_attachment=True,
            download_name='test_speech.wav'
        )
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    finally:
        # 稍后删除临时文件
        try:
            if temp_filename and os.path.exists(temp_filename):
                os.unlink(temp_filename)
        except:
            pass

if __name__ == '__main__':
    logger.info("===== 启动CosyVoice API服务 =====")
    
    if init_tts_system():
        # 运行测试
        logger.info("运行自测试...")
        test_text = "自测试文本，验证语音合成功能。"
        speech_np, sample_rate = generate_speech(test_text)
        if speech_np is not None:
            logger.info(f"自测试成功! 生成音频时长: {len(speech_np)/sample_rate:.2f}秒")
        else:
            logger.warning("自测试失败!")
        
        # 启动服务
        logger.info(f"服务启动: http://0.0.0.0:5001")
        app.run(host='0.0.0.0', port=5001, threaded=True)
    else:
        logger.error("无法启动API服务，模型加载失败")